{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6e8b2a6f",
   "metadata": {},
   "source": [
    "# Cohere reranker\n",
    "\n",
    ">[Cohere](https://cohere.ai/about)는 기업이 인간-기계 상호작용을 개선할 수 있도록 돕는 자연어 처리 모델을 제공하는 캐나다의 스타트업입니다.\n",
    "\n",
    "이 노트북은 `retriever`에서 [Cohere의 rerank endpoint](https://docs.cohere.com/docs/reranking)를 사용하는 방법을 보여줍니다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59e4104c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 설치\n",
    "# !pip install -qU cohere"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68f6dcf4",
   "metadata": {},
   "source": [
    "## Cohere API 키 설정\n",
    "\n",
    "- [API 키 발급](https://dashboard.cohere.com/api-keys)\n",
    "- `.env` 파일에 `COHERE_API_KEY` 키 값에 API 키를 넣어주세요.\n",
    "\n",
    "**참고**\n",
    "- [공식 도큐먼트](https://docs.cohere.com/docs/the-cohere-platform?_gl=1*st323v*_gcl_au*MTA2ODUyNDMyNy4xNzE4MDMzMjY2*_ga*NTYzNTI5NDYyLjE3MTgwMzMyNjc.*_ga_CRGS116RZS*MTcyMTk4NzMxMi4xMS4xLjE3MjE5ODczNjIuMTAuMC4w)\n",
    "- [Reranker Model 리스트](https://docs.cohere.com/docs/rerank-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b03e9b15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API KEY를 환경변수로 관리하기 위한 설정 파일\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# API KEY 정보로드\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f170a206",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangSmith 추적을 설정합니다. https://smith.langchain.com\n",
    "# !pip install langchain-teddynote\n",
    "from langchain_teddynote import logging\n",
    "\n",
    "# 프로젝트 이름을 입력합니다.\n",
    "logging.langsmith(\"Reranker\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b146db95",
   "metadata": {},
   "source": [
    "## 사용법"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efae0675",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pretty_print_docs(docs):\n",
    "    print(\n",
    "        f\"\\n{'-' * 100}\\n\".join(\n",
    "            [f\"Document {i+1}:\\n\\n\" + d.page_content for i, d in enumerate(docs)]\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae7d50e3",
   "metadata": {},
   "source": [
    "**Cohere 다국어 지원 모델**\n",
    "\n",
    "- Embedding: `embed-multilingual-v3.0`, `embed-multilingual-light-v3.0`, `embed-multilingual-v2.0`\n",
    "- Reranker: `rerank-multilingual-v3.0`, `rerank-multilingual-v2.0`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7a56ef1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import TextLoader\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from langchain_cohere import CohereEmbeddings\n",
    "\n",
    "# 문서 로드\n",
    "documents = TextLoader(\"./data/appendix-keywords.txt\").load()\n",
    "\n",
    "# 텍스트 분할기 초기화\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)\n",
    "\n",
    "# 문서 분할\n",
    "texts = text_splitter.split_documents(documents)\n",
    "\n",
    "# 검색기 초기화\n",
    "retriever = FAISS.from_documents(\n",
    "    texts, CohereEmbeddings(model=\"embed-multilingual-v3.0\")\n",
    ").as_retriever(search_kwargs={\"k\": 10})\n",
    "\n",
    "# 질의문\n",
    "query = \"Word2Vec 에 대해서 알려줘!\"\n",
    "\n",
    "# 문서 검색\n",
    "docs = retriever.invoke(query)\n",
    "\n",
    "# 문서 출력\n",
    "pretty_print_docs(docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82084994",
   "metadata": {},
   "source": [
    "## CohereRerank을 사용한 재정렬\n",
    "이제 기본 `retriever`를 `ContextualCompressionRetriever`로 감싸보겠습니다. Cohere 재정렬 엔드포인트를 사용하여 반환된 결과를 재정렬하는 `CohereRerank`를 추가할 것입니다.\n",
    "CohereRerank에서 모델 이름을 지정하는 것이 필수임을 유의하십시오!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b6d3d95",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers.contextual_compression import ContextualCompressionRetriever\n",
    "from langchain_cohere import CohereRerank\n",
    "\n",
    "# 문서 재정렬 모델 설정\n",
    "compressor = CohereRerank(model=\"rerank-multilingual-v3.0\")\n",
    "\n",
    "# 문맥 압축 검색기 설정\n",
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=compressor, base_retriever=retriever\n",
    ")\n",
    "\n",
    "# 압축된 문서 검색\n",
    "compressed_docs = compression_retriever.invoke(\"Word2Vec 에 대해서 알려줘!\")\n",
    "\n",
    "# 압축된 문서 출력\n",
    "pretty_print_docs(compressed_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a5457ba",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain-kr-lwwSZlnu-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
